import argparse
import json
import logging
import sys

import numpy as np
import redis
import zmq

import train
import util

# file path
CONFIG_PATH = f'{sys.path[0]}/config.json'

# database file
feature_db: np.ndarray
info_db: dict

# config folder
configs: dict
root_path: str

# redis cursor
rds: redis.StrictRedis


def insert(info: dict):
    global feature_db, info_db
    add_name_list = [name for name in list(info) if name not in set(list(info_db))]
    features, names = train.train_dataset(root_path, alter_paths=add_name_list)
    if features is None:
        logging.warning('[FR DB] invalid names!')
        return False
    # update features
    if feature_db is None:
        feature_db = np.array(features)
        info_db = info
    else:
        feature_db = np.append(feature_db, features, axis=0)
        info_db.update(info)
    util.to_redis_json(rds, info_db, configs['INFO_DB_KEY'])
    util.to_redis_numpy(rds, feature_db, configs['FEATURE_DB_KEY'])
    logging.debug(f'[FR DB] names inserted: {add_name_list}')
    return True


def delete(info: dict):
    global feature_db, info_db
    name_list = list(info_db)
    del_name_set = set(list(info))
    del_idx = [i + j for i in [name_list.index(name) * train.VALID_DATA_SIZE
                               for name in del_name_set.intersection(name_list)]
               for j in range(train.VALID_DATA_SIZE)]
    if len(del_idx) != 0:
        # update features
        feature_db = np.delete(feature_db, del_idx, 0)
        [info_db.pop(key) for key in info]
        util.to_redis_json(rds, info_db, configs['INFO_DB_KEY'])
        util.to_redis_numpy(rds, feature_db, configs['FEATURE_DB_KEY'])
        logging.debug(f'[FR DB] names deleted: {del_name_set}')
    else:
        logging.debug('[FR DB] all invalid names!')
        return False
    return True


def update(info: dict):
    info_db.update(info)
    util.to_redis_json(rds, info_db, configs['INFO_DB_KEY'])
    logging.debug(f'[FR DB] names updated: {list(info)}')
    return True


def truncate(values):
    global info_db, feature_db
    rds.flushdb()
    info_db.clear()
    feature_db = None
    logging.debug('[FR DB] db truncated')
    return True


def retrieve(values):
    global info_db, feature_db
    logging.info(f'[FR DB] current db: {info_db}')
    return False


def parse_sys_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--zmq_port', type=int, default=9000,
                        help='ZeroMQ publisher port of FRLite-DB.')
    parser.add_argument('-t', '--db_topic', type=str, default='db',
                        help='ZeroMQ publisher topic of FRLite-DB.')
    parser.add_argument('-l', '--log_level', type=str, choices=['DEBUG', 'INFO', 'WARN', 'ERROR', 'FATAL'],
                        help='Log level for showing logs.', default='INFO')
    return parser.parse_args(argv)


def init_input_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', type=str,
                        choices=['insert', 'delete', 'update', 'truncate', 'retrieve'],
                        help='Mode to use this file.')
    parser.add_argument('-v', '--values', type=json.loads, default='{"key": "value"}',
                        help='Values for mode (JSON String!).')
    return parser


def init(args):
    global configs, root_path, rds, feature_db, info_db
    logging.basicConfig(
        level=args.log_level,
        format='[%(asctime)s] <%(levelname)s> %(filename)s:%(lineno)d: %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S')
    # init param
    topic = args.db_topic
    # init jsons
    configs = dict(json.load(open(CONFIG_PATH, 'r', encoding='utf8')))
    root_path = configs['TRAIN_PATH']
    # init redis
    rds = redis.StrictRedis(host=configs['REDIS_HOST'],
                            port=configs['REDIS_PORT'],
                            db=configs['REDIS_DB'],
                            password=configs['REDIS_PASSWORD'])
    # init db
    info_db = dict(util.from_redis_json(rds, configs['INFO_DB_KEY']))
    feature_db = util.from_redis_numpy(rds, configs['FEATURE_DB_KEY'])
    logging.info(f'[FR DB] current db: {info_db}')
    # start zmq (db->lite)
    context = zmq.Context()
    socket: zmq.Socket = context.socket(zmq.PUB)
    socket.bind("tcp://0.0.0.0:{}".format(args.zmq_port))
    logging.info(f'[FR DB] ZeroMQ publisher started successfully at port {args.zmq_port}')
    # monitor inputs
    input_parser = init_input_parser()
    while True:
        # obtain new command
        i = input()
        args = input_parser.parse_args(i.split())
        # update db
        flag = eval("{}({})".format(str(args.mode), args.values))
        # broadcast update
        if flag:
            logging.info('[FR DB] waiting...')
            socket.send_multipart(
                [bytes(topic, 'utf-8'), bytes('DB', 'utf-8')])
            logging.info(f'[FR DB] current db: {info_db}')


if __name__ == '__main__':
    init(parse_sys_arguments(sys.argv[1:]))
